import math

import numpy as np

from agent.agent import Agent
from env.env_cmd import EnvCmd
from env.env_def import BLUE_AIRPORT_ID, UnitType

from .a2a import A2A
from .cmd import Command
from .task import Task, TaskState

TAKEOFF_PATROL_POINT = [0, 0, 8000]
TYPE2STRING_MAP = {11: "AIR", 12: "AWCS", 13: "JAM", 15: "BOM"}


class A2ASupportTask(Task):
    def __init__(self, task_state):
        super().__init__(task_state)
        self.done_unit_ids = []

    def _attack(self, obs_blue):
        cmds = []
        unit_ids = []
        en_unit_ids = []
        for _, unit in self.units_map.items():
            for en_unit in obs_blue["qb"]:
                if en_unit["LX"] in [11, 15]:
                    dist = unit.compute_2d_distance(en_unit["X"], en_unit["Y"])
                    if dist < 120000:
                        # cmds.append(Command.a2a_attack(unit.id, en_unit['ID']))
                        unit_ids.append(unit.id)
                        en_unit_ids.append(en_unit["ID"])
        # unit_ids = list(set(unit_ids))
        # en_unit_ids = list(set(en_unit_ids))
        # for unit_id in unit_ids:
        #     np.random.shuffle(en_unit_ids)
        #     for en_unit_id in en_unit_ids:
        #         cmds.append(Command.a2a_attack(unit_id, en_unit_id))
        return cmds


class A2ASupportPoint(A2ASupportTask):
    def __init__(self, max_unit_num, target_point, task_state=TaskState.PREPARE):
        super().__init__(task_state)
        self.max_unit_num = max_unit_num
        self.target_point = target_point

    def get_all_missiles(self):
        missile_num = 0
        for _, unit in self.units_map.items():
            missile_num += unit.get_missile_num()
        return missile_num

    def update_task_state(self):
        pass

    def update(self, alive_unit_map):
        self.update_units_map(alive_unit_map)

    def run(self, idle_unit_map, obs_blue):
        cmds = []
        attack_cmds = self._attack(obs_blue)
        cmds.extend(attack_cmds)
        if self.task_state != TaskState.FINISH:
            while len(self.units_map) < self.max_unit_num:
                if len(idle_unit_map) == 0:
                    break
                unit = idle_unit_map.popitem()[1]
                self.add_unit(unit)
            if len(attack_cmds) == 0:
                for unit_id, _ in self.units_map.items():
                    cmds.append(
                        Command.area_patrol(
                            unit_id, [self.target_point[0], self.target_point[1], 8000]
                        )
                    )
        else:
            for unit_id, _ in self.units_map.items():
                cmds.append(Command.area_patrol(unit_id, TAKEOFF_PATROL_POINT))
            self.finish()
        return cmds


class BlueAgent(Agent):
    def __init__(self, name, config, **kwargs):
        super().__init__(name, config["side"])
        self.a2a_task = A2ASupportPoint(3, [0, 0])
        self.find = False
        self.num = 0

    def _get_waiting_aircraft_num(self, obs_blue, type_):
        if type_ not in TYPE2STRING_MAP.keys():
            return 0
        type_str = TYPE2STRING_MAP[type_]
        return obs_blue["airports"][0][type_str]

    """
        控制飞机起飞
    """

    def _take_off(self, obs_blue):
        cmds = []
        # fly_types = [UnitType.A2A, UnitType.A2G, UnitType.JAM]
        fly_types = [UnitType.A2A]
        for type_ in fly_types:
            if self._get_waiting_aircraft_num(obs_blue, type_):
                cmds.append(
                    Command.takeoff_areapatrol(
                        BLUE_AIRPORT_ID,
                        1,
                        type_,
                        patrol_points=[0, 0 - self.num * 100000, 8000],
                    )
                )
                self.num += 1

        return cmds

    def _awacs_task(self, obs_blue):
        cmds = []
        patrol_points = [-180000, 150000, 8000]
        for unit in obs_blue["units"]:
            if unit["LX"] == UnitType.AWACS:
                cmds.append(Command.awacs_areapatrol(unit["ID"], patrol_points))

                break
        return cmds

    def _get_units_map(self, obs_blue, type_):
        units_map = {}
        for info_map in obs_blue["units"]:
            if info_map["LX"] == type_:
                units_map[info_map["ID"]] = A2A(info_map)
        return units_map

    def step(self, sim_time, obs, **kwargs):
        cmds = []
        obs_red = obs["red"]
        obs_blue = obs["blue"]
        self.find = True
        for unit in obs_blue["units"]:
            if unit["LX"] == UnitType.AWACS:
                cmds.append(EnvCmd.make_awcs_radarcontrol(unit["ID"], 1))
            # if unit["LX"] == UnitType.A2A:
            #     self.find = False
        if obs_blue["airports"][0]["AIR"] > 0 and sim_time > 0:
            cmds.extend(self._take_off(obs_blue))
            # cmds.extend(self._awacs_task(obs_blue))  # 预警机巡逻
            return cmds

        # a2a_map = self._get_units_map(obs_blue, UnitType.A2A)
        # self.a2a_task.update(a2a_map)
        # for unit_id, _ in self.a2a_task.get_units_map().items():
        #     a2a_map.pop(unit_id)
        # cmds.extend(self.a2a_task.run(a2a_map, obs_blue))

        if self.find:
            # obs是观测的
            self_a2a = []
            enemy_a2a = []
            my_a2a = 0
            # 遍历蓝方的units，看是否是飞机
            for unit in obs_blue["units"]:
                if unit["LX"] == UnitType.A2A:
                    my_a2a += 1
                    self_a2a.append(unit)
            for unit in obs_red["units"]:
                if unit["LX"] == UnitType.A2A:  # 18：未知空中目标
                    enemy_a2a.append(unit)
            for i in range(my_a2a):
                i_unit = self_a2a[i]
                # 提取坐标
                x1, y1 = i_unit["X"], i_unit["Y"]
                distance0 = 20000
                angle_blue = 0
                for unit in enemy_a2a:
                    x2, y2 = unit["X"], unit["Y"]
                    # 计算向量
                    vector = np.array([x2 - x1, y2 - y1])

                    distance = np.linalg.norm(vector)
                    if distance < distance0:
                        distance0 = distance
                        # 计算角度（弧度）
                        angle_blue = np.arctan2(vector[1], vector[0])
                if distance0 < 20000:
                    cmds.append(
                        EnvCmd.make_areapatrol(
                            i_unit["ID"],
                            i_unit["X"] - 5000 * np.cos(angle_blue).item(),
                            i_unit["Y"] - 5000 * np.sin(angle_blue).item(),
                            i_unit["Z"],
                            0,
                            1000,
                            1000,
                            min(distance0 / 40, 150),
                            7200,
                        )
                    )
            # 速度稍慢于红方，位置是往敌方飞机的反方向
        return cmds

    def reset(self):
        self.find = False
        self.num = 0
        pass

    def get_distance(self, unit1, unit2):
        d_x = unit1["X"] - unit2["X"]
        d_y = unit1["Y"] - unit2["Y"]
        return math.sqrt(d_x**2 + d_y**2)
